# Sourabh Jain 
# sourj@cs.umn.edu

# This file uses 'hmetis' to perform the vid construction

# the input file is the node graph file. 
# ndges nnodes
# edge1
# edge2

#python assing_binary_ids_using_hmetis.py tmp/4755.r1.graph 4755.r1.vid


import sys, commands, time

if len(sys.argv) != 3:
    print 'Usage: ', sys.argv[0], ' <i/p .graph file> <o/p vid file>'
    sys.exit()
#----------------------------------------------------
#			FIXIFNOTCONNECTED START:
#----------------------------------------------------
def fixIfNotConnected(mapping, edges, parts):
    print 'fixIfNotConnected is called!'
    connected = 0
    # first see if the input graph is connected or not?
    groupids = {}
    for n in mapping:
        groupids[n] = n
    for ed in edges:
        if groupids[ed[0]] != groupids[ed[1]]:
            newlabel = min (groupids[ed[0]],groupids[ed[1]])
            oldlabel = max (groupids[ed[0]],groupids[ed[1]])
            for n in groupids:
                if groupids[n] == oldlabel:
                    groupids[n] = newlabel
    groupcounter = {}
    for n in groupids:
        if groupids[n] not in groupcounter:
            groupcounter[groupids[n]] = 0
        groupcounter[groupids[n]] = groupcounter[groupids[n]] + 1
    if len(groupcounter) > 1:
        print 'the input graph is disconnected!'
        print groupcounter
        sys.exit()
    groupids = {}
    groupcounter = {}
    while connected == 0:
		mapping_labels = []
		for m in mapping:
			#print mapping[m]
			if mapping[m] not in mapping_labels:
				mapping_labels.append(mapping[m])
		#print '--> Mapping labels are: ', mapping_labels
		labels = {}
		for n in mapping:
			labels[n] = n
		#print '---> Parts are : \n', parts
		for part in parts:
			for ed in part:
				if labels[ed[0]] != labels[ed[1]]:
					#merge these two groups
					newlabel  = min (labels[ed[0]], labels[ed[1]])
					oldlabel  = max (labels[ed[0]], labels[ed[1]])
					for n in labels:
						if labels[n] == oldlabel:
							labels[n] = newlabel
		#print '---> labels are \n', labels
		#Now we have number of disconnected groups:
		#Next step is to merge the smallest one with the other group
		groups= {}
		for n in labels:
			if labels[n] not in groups:
				groups[labels[n]] = 0
			groups[labels[n]] = groups[labels[n]]+1
		#see if number of such groups are more than 2
		print '---> Groups are: \n',groups
		if len(groups) == 2:
			# if the number of groups are 2 than no more fixing is required
			connected = 1
			break
		if len(groups) <2:
			# something is definitely wrong
			# plot the groups and stop
			plot (edges, mapping)
			print 'Number of groups are less than 2,[', len(groups), ']! Stopping now!'
			sys.exit()
		
		#Now we have more than 2 groups, therefore we'll require merging of the smallest with the other part.
		
		#find the smallest:
		smallestgrouplabel = -1
		sizeofsmallestgroup = -1
		for g in groups:
			#print 'Current smallest group is: ', smallestgrouplabel,' size is: ', sizeofsmallestgroup
			if sizeofsmallestgroup == -1 or sizeofsmallestgroup > groups[g]:
				#print 'Current smallest group is: ', smallestgrouplabel,' size is: ', sizeofsmallestgroup, ' and the new smallest group size is: ', groups[g], ' and group label is : ', g
				sizeofsmallestgroup = groups[g]
				smallestgrouplabel = g
		print 'Smallest group label is ', smallestgrouplabel, ' and the size is ', sizeofsmallestgroup
		# now for the nodes in the smallest group switch the mappings:
		for n in mapping:
			if labels[n] == smallestgrouplabel:
				if mapping[n] == mapping_labels[0]:
					mapping[n] = mapping_labels[1]
				else:
					mapping[n] = mapping_labels[0]
		
		# now update the parts
		parts = [[],[]]
		for edge in edges:
			if mapping[edge[0]] == mapping_labels[0] and mapping[edge[1]] == mapping_labels[0]:
				#then this edge belongs to part0
				parts[0].append(edge)
			elif mapping[edge[0]] != mapping_labels[0] and mapping[edge[1]] != mapping_labels[0]:
				#then this edge belongs to part0
				parts[1].append(edge)
		#plot the udpated graph
		plot(edges, mapping)
    return parts
	
#----------------------------------------------------
#			FIXIFNOTCONNECTED ENDS:
#----------------------------------------------------    
#----------------------------------------------------
#			GETNODES START:
#----------------------------------------------------
def getnodes(edges):
	nodes = {}
	for e in edges:
		nodes[e[0]] = 0
		nodes[e[1]] = 0
	return nodes
#----------------------------------------------------
#			GETNODES ENDS:
#----------------------------------------------------
#----------------------------------------------------
#			delta START:
#----------------------------------------------------
def delta(id1, id2):
	l = len(id1)
	r = 0
	for i in range(0, len(id1)):
		if id1[i] == id2[i]:
			r = r+1
		else:
			break
	return (l-r)
#----------------------------------------------------
#			delta ENDS:
#----------------------------------------------------
def plot(part0, mappings):
    #if len(mappings) > 500:
       # print 'Am not printing it, too large a graph it is! size = ', len(mappings)
        #return
    
	ofile = open('./tmp/graph.dot', 'w')
	ofile.write('graph G{')
	color1 = 'red'
	color2 = 'green'
	for n in mappings:
		color = color2
		if mappings[n] == '0':
			color = color1

		ofile.write(str(n) + ' [color='+color+'];\n')
	for e in part0:
		if e[0] < e[1]:
			ofile.write(str(e[0]) + '--'+str(e[1])+';\n')
	ofile.write('}')
	ofile.close()
	print commands.getoutput('dot ./tmp/graph.dot -Tpng -o ./tmp/graph_partitioning.'+str(int(time.time()*100)) + '.png')  
#----------------------------------------------------
#		VID_ASSIGNMENT START:
#----------------------------------------------------
# pseudocode:
# get the bi-partitioning first
# 	Now assign the first half a 0 and second half a 1
# 	split the vertices set 
# def: vid_assignment(part0, currentvids)
#	if (part0 > 1):
# 	part1, part2 = bi-partition(part0);
# 	idappend(part1, 0, currentvids)
# 	idappend(part2, 1, currentvids)
# 	currentvids = vid_assignment(part1, currentvids)
#	currentvids = vid_assignment(part2,currentvids)
def vid_assignment(part0, currentvids):
	if len(part0) == 0:
		print 'Empty edges!! \n'
		return currentvids
	nnodes = 0
	nedges = len(part0)
	nodes2newid = {}
	newid2nodes = {}

	for edge in part0:
		v1 = edge[0]
		v2 = edge[1]
		if v1 not in nodes2newid:
			nnodes = nnodes + 1
			nodes2newid[v1] = nnodes
			newid2nodes[nnodes] = v1
		if v2 not in nodes2newid:
			nnodes = nnodes + 1
			nodes2newid[v2] = nnodes
			newid2nodes[nnodes] = v2
	nnodes = len(nodes2newid)
	#print 'EDGES:::\n',edges
	#print 'NODES2NEWID:::\n',nodes2newid
	#print 'NEWID2NODES:::\n',newid2nodes
	# if nnodes < 2: then exit
	if nnodes < 3:
		for edge in part0:
			currentvids[edge[0]] = currentvids[edge[0]] + '0'
			currentvids[edge[1]] = currentvids[edge[1]] + '1'
			#print edge[0],currentvids[edge[0]]
			#print edge[1], currentvids[edge[1]]
		print 'Reached the leaves!'
		return currentvids

	ofile = open('./tmp/graph.' +str(nnodes), 'w')
	ofile.write(str(len(part0)) + ' '+ str(nnodes) + '\n')
	for edge in part0:
		ofile.write(str(nodes2newid[edge[0]]) + ' ' + str(nodes2newid[edge[1]]) + '\n')
	ofile.close()
	
	# now run hmetis to get two sub-graphs
	output =  commands.getoutput('shmetis ./tmp/graph.'+str(nnodes) + ' 2 1')
	#print output
	output = commands.getoutput('cat ./tmp/graph.'+str(nnodes) + '.part.2')
	mapping = output.split('\n')
	# now read the output to get the part1 and part 2
	
	node2partmapping = {}
	#print newid2nodes
	#print mapping
	groups = {}
	for i in range (0, len(mapping)):
		#print i,'th-->\t',mapping[i]
		groups[mapping[i]] = 1
		node2partmapping [newid2nodes[i+1]] = mapping[i]
	nodes = nodes2newid.keys()
	nodes.sort()
	print '# of groups are: ',len(groups)#, '\nAnd the nodes are: ',nodes
	parts = [[],[]]
	#print 'NODE2PARTMAPPING::\n', node2partmapping
	# now prepare the new parts
	group1 = node2partmapping[newid2nodes[1]]
	#print 'GROUP1 ID:::\n', group1
	for edge in part0:
		if node2partmapping[edge[0]] == group1 and node2partmapping[edge[1]] == group1:
			#then this edge belongs to part0
			parts[0].append(edge)
		elif node2partmapping[edge[0]] != group1 and node2partmapping[edge[1]] != group1:
			#then this edge belongs to part0
			parts[1].append(edge)
	plot(part0, node2partmapping)
	
	# if there are some disconnected parts then fix them
	parts = fixIfNotConnected(node2partmapping, part0, parts)
	#Now update the currenvids
	for node in node2partmapping:
		if node2partmapping[node] == group1:
			currentvids[node] = currentvids[node] +'0'
		else:
			currentvids[node] = currentvids[node] +'1'				
	currentvids = vid_assignment(parts[0], currentvids)
	currentvids = vid_assignment(parts[1], currentvids)
	#Now makesure that vids for both parts are of same length:
	#To achieve this, we will compare the difference in the length of total bits used
	# in both parts, and append 0 at the end of the part which has smaller number of bits used
	group1 = node2partmapping[newid2nodes[1]]
	len1 = len(currentvids[newid2nodes[1]])
	len2 = len1
	#print newid2nodes
	#print node2partmapping
	#print group1
	for node in node2partmapping:
		if node2partmapping[node] != group1:
			len2 = len(currentvids[node])
	diff = len1-len2
	print 'diff: ',diff
	#print 'part 0', parts[0]
	#print 'part 1',parts[1]
	#print currentvids
	#diff = 0
	if diff != 0:
		if diff > 0:
			ntoappend = diff
			#append ntoappend 0 to group2
			string2append = ''
			for i in range (0, ntoappend):
				string2append = string2append + '0'
			for node in node2partmapping:
				if node2partmapping[node] != group1:
					currentvids[node] = currentvids[node] + string2append
		else:
			ntoappend = -diff
			#append ntoappend 0 to group2
			string2append = ''
			for i in range (0, ntoappend):
				string2append = string2append + '0'
			for node in node2partmapping:
				if node2partmapping[node] == group1:
					currentvids[node] = currentvids[node] + string2append					
	#print node2partmapping.keys()				
	#for node in node2partmapping:
		#print 'Vid for node: ', node, ': ',currentvids[node]
	return currentvids
#----------------------------------------------------
#		VID_ASSIGNMENT START:
#----------------------------------------------------


#----------------------------------------------------
#		MAIN METHOD START:
#----------------------------------------------------
# construct the edge set first:
infile = open(sys.argv[1], 'r')
line = infile.readline()
line = infile.readline()
currentvids = {}
edges = []
# remove old plots and graphs.
print commands.getoutput('rm -rf tmp/graph_partitioning*.png; rm -rf tmp/graph.*')
while line != '':
	line = line.strip()
	v = line.split(' ')
	v1 = int(v[0])
	v2 = int (v[1])
	currentvids[v1] = ''
	currentvids[v2] = ''
	if v1 < v2:
		edge = [v1,v2]
		edges.append(edge)
	line = infile.readline()
infile.close()

currentvids = vid_assignment(edges, currentvids)
# now write down the mappings:
outfile = open(sys.argv[2], 'w')
#print currentvids
for node in currentvids:
	outfile.write(str(node) + ' ' + str(int(currentvids[node],2)) + '\n')
outfile.close()

outfile = open(sys.argv[2]+'.bin', 'w')
#print currentvids
for node in currentvids:
    outfile.write(str(node) + ' "' + currentvids[node] + '"\n')
outfile.close()
# Create a graph figure.
outfile = open(sys.argv[2] + '.dot','w')
outfile.write('graph G{\n')
for node in currentvids:
	outfile.write(str(node) + ' [label='+ str(int(currentvids[node],2))+'];\n')
for edge in edges:
	if edge[0] < edge[1]:
		edge1 = str(edge[0]) + '--' + str(edge[1])+' [label='+str(delta(currentvids[edge[0]], currentvids[edge[1]])) + ']'
    	outfile.write(edge1+';\n')
outfile.write('}\n')
outfile.close()
print commands.getoutput('dot '+sys.argv[2]+'.dot -Tpng -o '+sys.argv[2]+'.png')

